package net.freebytes.es.demo.service;

import net.freebytes.es.common.repository.*;
import org.elasticsearch.index.query.QueryBuilder;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.util.ArrayList;
import java.util.List;

@Service
public class BaseQueryService {

    @Autowired
    private ESTemplate esTemplate;

    private final int defaultPageSize = 10;

    /**
     * 以search_after的方式 一次性获取全量数据
     *
     * @author 千里明月
     * @date 2021/11/17
     **/
    public <T> List<T> searchAfter(QueryBuilder queryBuilder, Sort sort, Class<T> clazz, String[] includes) {
        int currentPage = 1;
        int pageSize = defaultPageSize;
        PageSortHighLight psh = new PageSortHighLight(currentPage, pageSize);
        if (sort != null) {
            psh.setSort(sort);
        }
        Attach attach = new Attach();
        attach.setSearchAfter(true);
        attach.setPageSortHighLight(psh);
        if (includes != null && includes.length > 0) {
            attach.setIncludes(includes);
        }
        List<T> list = new ArrayList<T>();
        // 获取第一页数据
        PageList<T> page = esTemplate.search(queryBuilder, attach, clazz);

        List<T> queryList = page.getList();

        //分别获取每页数据
        while (queryList != null && queryList.size() > 0) {
            list.addAll(queryList);
            // search_after循环请求
            queryList.clear();
            queryList = null;
            attach.setSortValues(page.getSortValues());
            page = esTemplate.search(queryBuilder, attach, clazz);
            queryList = page.getList();
        }
        return list;
    }

    public <T> PageList<T> searchAfterOnePage(QueryBuilder queryBuilder, Class<T> clazz, PageList<T> page) {
        int currentPage = 1;
        int pageSize = 10;
        PageSortHighLight psh = new PageSortHighLight(currentPage, pageSize);

        Attach attach = new Attach();
        attach.setSearchAfter(true);
        attach.setPageSortHighLight(psh);
        if (page != null) {
            attach.setSortValues(page.getSortValues());
        }
        // 获取一页数据
        return esTemplate.search(queryBuilder, attach, clazz);
    }
}
